import argparse

from yolov3.model.darknet import load_model
from yolov3.utils.parse_config import parse_data_config
from yolov3.eval.eval import evaluate, create_validation_data_loader
from yolov3.utils.utils import load_classes
from yolov3.utils.env import print_environment_info


def run(args):
    print_environment_info()
    # Load configuration from data file
    data_config = parse_data_config(args.data)
    # Path to file containing all images for validation
    valid_path = data_config["valid"]
    class_names = load_classes(data_config["names"])  # List of class names

    dataloader = create_validation_data_loader(
        valid_path, args.batch_size, args.img_size, args.n_cpu
    )
    model = load_model(args.model, args.device, args.weights)
    metrics_output = evaluate(
        model,
        args.device,
        dataloader,
        class_names,
        args.img_size,
        args.iou_thres,
        args.conf_thres,
        args.nms_thres,
        verbose=True,
    )


def parse_args():
    parser = argparse.ArgumentParser(description="Evaluate validation data.")

    # fmt: off
    parser.add_argument("-m", "--model", type=str, default="config/yolov3.cfg", help="Path to model definition file (.cfg)")
    parser.add_argument("-d", "--device", type=str, default="cuda:0", help="Device used for detection")
    parser.add_argument("-w", "--weights", type=str, default="weights/yolov3.weights", help="Path to weights or checkpoint file (.weights or .pth)")
    parser.add_argument("-d", "--data", type=str, default="config/coco.data", help="Path to data config file (.data)")
    parser.add_argument("-b", "--batch-size", type=int, default=8, help="Size of each image batch")
    parser.add_argument("-v", "--verbose", action='store_true', help="Makes the validation more verbose")
    parser.add_argument("--img-size", type=int, default=416, help="Size of each image dimension for yolo")
    parser.add_argument("--n-cpu", type=int, default=8, help="Number of cpu threads to use during batch generation")
    parser.add_argument("--iou-thres", type=float, default=0.5, help="IOU threshold required to qualify as detected")
    parser.add_argument("--conf-thres", type=float, default=0.01, help="Object confidence threshold")
    parser.add_argument("--nms-thres", type=float, default=0.4, help="IOU threshold for non-maximum suppression")
    # fmt: on

    args = parser.parse_args()
    print(f"Command line arguments: {args}")

    return args


if __name__ == "__main__":
    args = parse_args()
    run(args)
